#ifndef HAARWAVELET_HEADER
#define HAARWAVELET_HEADER

#include <glm\glm.hpp>
#include <vector>

const float SQRT2 = std::sqrt(2.0f);

struct HaarDecomp {

	HaarDecomp() :
numLevs(0), totalDim(0), coeffDim(0) {}
	
	int GetDimLevel( int level ) {
		if (level > numLevs){
			fprintf(stderr, "Error: %d is greater than numLevs %d\n", level, numLevs);
			return 0;
		}

		return std::pow(2.0, numLevs - level ) * coeffDim;
	}

	int numLevs;
	int totalDim;
	int coeffDim;

	std::vector<glm::vec2> mData;
};

class HaarWavelet {
public:


	void SetData( const std::vector<glm::vec2> & data, int dim ){
		mData = data;
		mDim = dim;
	}

	std::vector<glm::vec2> * GetData(){
		return &mData;
	}

	void Decompose(int numLevs = -1) {
	
		mDecomp = HaarDecomp();
		mDecomp.mData.resize(mDim*mDim);
		mDecomp.coeffDim = mDim;
		mDecomp.totalDim = mDim;


		// Normalize
		for (int i = 0; i < mDecomp.mData.size(); ++i){
			 mDecomp.mData[i] = mData[i] / (float)mDim;
		}


		Decompose(mDecomp, numLevs);
	}

	// Decompose based on the input decomposition.
	// Will continue from the current amount of decomposition.
	void Decompose(HaarDecomp & inDecomp, int numLevs = -1){

		int maxLevs = (int)(std::log((double)inDecomp.totalDim) / std::log(2.0));
		int decompLevs = maxLevs - mDecomp.numLevs;

		if (numLevs == -1 || numLevs > decompLevs){
			mDecomp.numLevs = maxLevs;
		}
		else {
			decompLevs = numLevs;
			mDecomp.numLevs += decompLevs;
		}


		int dim = inDecomp.totalDim;
		int g = inDecomp.coeffDim;

		int cntr = 0;
		while (g >= 2 && cntr < decompLevs){

			for (int r = 0; r < g; ++r){
				Decompose1D(&inDecomp.mData[r*dim], 1, g);
			}
			for (int c = 0; c < g; ++c){
				Decompose1D(&inDecomp.mData[c], dim, g);
			}
			g = g/2;
			inDecomp.coeffDim = g;
			cntr++;
		}
	}



	void Decompose1D( glm::vec2 * in, int stride, int num ) {

		std::vector<glm::vec2> tmp(num);

		for (int i = 0; i < num/2; ++i){
			glm::vec2 c0 = in[(2*i)*stride],
				c1 = in[(2*i+1)*stride];

			tmp[i] = (c0 + c1) / SQRT2;
			tmp[(num/2 + i)] = (c0 - c1) / SQRT2;
		}

		for (int i = 0; i < num; ++i){
			in[i*stride] = tmp[i];
		}
	}

	void Reconstruct1D(glm::vec2 * in, int stride, int num){
		std::vector<glm::vec2> tmp(num);

		for (int i = 0; i < num/2; ++i){
			glm::vec2 c0 = in[i*stride],
				c1 = in[(num/2 + i)*stride];

			tmp[(2*i)] = (c0 + c1) / SQRT2;
			tmp[(2*i+1)] = (c0 - c1) / SQRT2;

		}

		for (int i = 0; i < num; ++i){
			in[i*stride]= tmp[i];
		}
	}

	void Reconstruct() {

		mData = mDecomp.mData;

		int g = mDecomp.coeffDim*2;
		while (g <= mDim){

			for (int c = 0; c < g; ++c){
				Reconstruct1D(&mData[c], mDim, g);
			}

			for (int r = 0; r < g; ++r){
				Reconstruct1D(&mData[r*mDim], 1, g);
			}

			g = 2*g;
		}

		for (int i = 0; i < mData.size(); ++i){
			mData[i]*=mDim;
		}

	}

	HaarDecomp * GetDecomp() {
		return &mDecomp;
	}


private:
	int mDim;
	std::vector<glm::vec2> mData;

	HaarDecomp mDecomp;

};

#endif